{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "cZGVhz8wgtqi"
      },
      "source": [
        "# Multimodal NER Tutorial\n",
        "> Tutorial author:乔硕斐（shuofei@zju.edu.cn）\n",
        "\n",
        "In this tutorial, we use a Transformer based two-stream multimodal model——IFAformer to recognize named entities. We hope this tutorial can help you understand the process of multimodal named entity recognition.\n",
        "\n",
        "This tutorial uses `Python3`."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "GwsLWjt3jv5B"
      },
      "source": [
        "## NER\n",
        "**Named-entity recognition** (also known as named entity identification, entity chunking, and entity extraction) is a subtask of information extraction that seeks to locate and classify named entities mentioned in unstructured text into pre-defined categories such as person names, organizations, locations, medical codes, time expressions, quantities, monetary values, percentages, etc."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "TGWSZX6_kgzy"
      },
      "source": [
        "## Multimodal NER\n",
        "**Multimodal named entity recognition**(MNER) in DeepKE applies a correlational visual image for each sentence to enhance the textual named entity recognition."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "yiv37DAHlOQx"
      },
      "source": [
        "## Dataset\n",
        "We use Twitter15 dataset from [UMT](https://github.com/jefferyYu/UMT/) in this tutorial. Each piece of data contains a sentence and an image. The data formats are as follow:\n",
        "\n",
        "**Text**\n",
        "\n",
        "```\n",
        "IMGID:16_05_01_6\n",
        "5\tO\n",
        "days\tO\n",
        "until\tO\n",
        "JUSTIN\tB-PER\n",
        "BIEBER\tI-PER\n",
        "CONCERT\tO\n",
        "omfg\tO\n",
        "Ima\tO\n",
        "die\tO\n",
        "#\tO\n",
        "puposetour\tO\n",
        "@\tO\n",
        "alminababexox\tO\n",
        "@\tO\n",
        "justinbieber\tB-PER\n",
        "```\n",
        "\n",
        "**Image**\n",
        "\n",
        "![image](https://github.com/zjunlp/DeepKE/blob/main/tutorial-notebooks/ner/multimodal/image/image.jpg?raw=1)\n",
        "\n",
        "The structure of the dataset folder `./data/` is as follow:\n",
        "\n",
        "```\n",
        ".\n",
        "├── twitter15_detect                     # Detected objects using RCNN\n",
        "├── twitter2015_aux_images                  # Detected objects using visual grounding toolkit\n",
        "├── twitter2015_images                    # Original images\n",
        "├── train.txt                        # Train set\n",
        "└── ...\n",
        "```\n",
        "\n",
        "We use RCNN detected objects and visual grounding objects as visual local information, where RCNN via [faster_rcnn](https://github.com/pytorch/vision/blob/main/torchvision/models/detection/faster_rcnn.py) and visual grounding via [onestage_grounding](https://github.com/zyang-ur/onestage_grounding)."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "FtHGsVPwtOgz"
      },
      "source": [
        "## IFAformer\n",
        "IFAformer is a novel dual Multimodal Transformer model with implicit feature alignment for the RE task, which utilizes the Transformer structure uniformly in the visual and textual without explicitly designing modal alignment structure. Here **we add a CRF module instead of the relation predictor** after IFAformer to enforce NER task.\n",
        "\n",
        "![IFAformer](https://github.com/zjunlp/DeepKE/blob/main/tutorial-notebooks/ner/multimodal/image/ner.png?raw=1)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "VFRjdJhLuEyZ"
      },
      "source": [
        "## Prepare environment"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 1000
        },
        "id": "bD2-vlJjuHrV",
        "outputId": "4daa25bd-4085-4bbe-a9e4-6efadec1d67c"
      },
      "outputs": [],
      "source": [
        "! nvidia-smi\n",
        "! pip install deepke\n",
        "! git clone https://github.com/zjunlp/DeepKE.git"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "0OwYr8kAuZ0q"
      },
      "source": [
        "Since the **torchvision** provided by Colab is the latest version, it is incompatible with the **torch** version required by deepke, so it needs to be unified. **This step can be ignored in non Colab environments.**"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "TB5i8II5uck8",
        "outputId": "1aae8553-fe73-4bd8-c1d6-563030407bec"
      },
      "outputs": [],
      "source": [
        "! pip install torch==1.10.0+cu102 torchvision==0.11.0+cu102 torchaudio==0.10.0 -f https://download.pytorch.org/whl/torch_stable.html"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "UEHVBwpauhoB"
      },
      "source": [
        "## Import package"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "kjAOtGKEu4kR"
      },
      "source": [
        "The python environment provided by Colab is 3.7, which is inconsistent with the python=3.8 required by deepke. To solve the problem, it is needed to fix ***import importlib.metadata as importlib\\_metadata*** to ***import importlib\\_metadata*** in file /usr/local/lib/Python3.7/dist-packages/deepke/relation\\_extraction/multimodal/models/clip/file\\_utils.py. and /usr/local/lib/Python3.7/dist-packages/deepke/name\\_entity\\_re/multimodal/models/clip/file\\_utils.py. **For non colab environments, please use python3.8, and ignore this step.**"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "1Z2GiUmMu5L7"
      },
      "outputs": [],
      "source": [
        "import os\n",
        "import hydra\n",
        "import torch\n",
        "import numpy as np\n",
        "import random\n",
        "from PIL import Image\n",
        "from torch.utils.data import Dataset\n",
        "from transformers import BertTokenizer\n",
        "from hydra import utils\n",
        "from deepke.name_entity_re.multimodal.models.clip.processing_clip import CLIPProcessor\n",
        "from torch.utils.data import DataLoader\n",
        "from deepke.name_entity_re.multimodal.models.IFA_model import IFANERCRFModel\n",
        "from deepke.name_entity_re.multimodal.modules.dataset import MMPNERProcessor, MMPNERDataset\n",
        "from deepke.name_entity_re.multimodal.modules.train import Trainer\n",
        "\n",
        "import warnings\n",
        "warnings.filterwarnings(\"ignore\", category=UserWarning)\n",
        "import logging\n",
        "logging.basicConfig(format = '%(asctime)s - %(levelname)s - %(name)s -   %(message)s',\n",
        "                    datefmt = '%m/%d/%Y %H:%M:%S',\n",
        "                    level = logging.INFO)\n",
        "logger = logging.getLogger(__name__)\n",
        "\n",
        "# import wandb\n",
        "# writer = wandb.init(project=\"DeepKE_NER_MM\")\n",
        "writer=None"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "E8yIeZChvGAj"
      },
      "source": [
        "## Prepare dataset"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "P5M0T1xrvHf3"
      },
      "outputs": [],
      "source": [
        "! wget 120.27.214.45/Data/ner/multimodal/data.tar.gz\n",
        "! tar -xzvf data.tar.gz"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "uKtrCHYQvR77"
      },
      "outputs": [],
      "source": [
        "DATA_PATH = {\n",
        "    'twitter15': {'train': 'data/twitter2015/train.txt',\n",
        "                'dev': 'data/twitter2015/valid.txt',\n",
        "                'test': 'data/twitter2015/test.txt',\n",
        "                'train_auximgs': 'data/twitter2015/twitter2015_train_dict.pth',\n",
        "                'dev_auximgs': 'data/twitter2015/twitter2015_val_dict.pth',\n",
        "                'test_auximgs': 'data/twitter2015/twitter2015_test_dict.pth',\n",
        "                'rcnn_img_path': 'data/twitter2015',\n",
        "                'img2crop': 'data/twitter2015/twitter2015_img2crop.pth'},\n",
        "\n",
        "    'twitter17': {'train': 'data/twitter2017/train.txt',\n",
        "                'dev': 'data/twitter2017/valid.txt',\n",
        "                'test': 'data/twitter2017/test.txt',\n",
        "                'train_auximgs': 'data/twitter2017/twitter2017_train_dict.pth',\n",
        "                'dev_auximgs': 'data/twitter2017/twitter2017_val_dict.pth',\n",
        "                'test_auximgs': 'data/twitter2017/twitter2017_test_dict.pth',\n",
        "                'rcnn_img_path': 'data/twitter2017',\n",
        "                'img2crop': 'data/twitter2017/twitter17_img2crop.pth'}\n",
        "    }\n",
        "\n",
        "IMG_PATH = {\n",
        "    'twitter15': 'data/twitter2015/twitter2015_images',\n",
        "    'twitter17': 'data/twitter2017/twitter2017_images'\n",
        "}\n",
        "\n",
        "AUX_PATH = {\n",
        "    'twitter15': {'train': 'data/twitter2015/twitter2015_aux_images/train/crops',\n",
        "                'dev': 'data/twitter2015/twitter2015_aux_images/val/crops',\n",
        "                'test': 'data/twitter2015/twitter2015_aux_images/test/crops'},\n",
        "\n",
        "    'twitter17': {'train': 'data/twitter2017/twitter2017_aux_images/train/crops',\n",
        "                'dev': 'data/twitter2017/twitter2017_aux_images/val/crops',\n",
        "                'test': 'data/twitter2017/twitter2017_aux_images/test/crops'}\n",
        "}\n",
        "\n",
        "LABEL_LIST = [\"O\", \"B-MISC\", \"I-MISC\", \"B-PER\", \"I-PER\", \"B-ORG\", \"I-ORG\", \"B-LOC\", \"I-LOC\", \"X\", \"[CLS]\", \"[SEP]\"]"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "QEijWUo1vdiy"
      },
      "source": [
        "## Configure parameters"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "z0SANyZDvfth"
      },
      "outputs": [],
      "source": [
        "class Config(object):\n",
        "    seed = 1234\n",
        "    \n",
        "    bert_name = \"bert-base-uncased\"\n",
        "    vit_name = \"openai/clip-vit-base-patch32\"\n",
        "    device = \"cuda\"\n",
        "    \n",
        "    num_epochs = 30\n",
        "    batch_size = 32\n",
        "    lr = 5e-5\n",
        "    warmup_ratio = 0.06\n",
        "    eval_begin_epoch = 1\n",
        "\n",
        "    max_seq = 40\n",
        "    aux_size = 128\n",
        "    rcnn_size = 64\n",
        "    ignore_idx = 0\n",
        "\n",
        "    save_path = \"checkpoints/twitter15\"\n",
        "    load_path = None\n",
        "\n",
        "    dataset_name = \"twitter15\"\n",
        "\n",
        "cfg = Config()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "sO2obEdBv3qA"
      },
      "outputs": [],
      "source": [
        "def set_seed(seed=2021):\n",
        "    \"\"\"set random seed\"\"\"\n",
        "    torch.manual_seed(seed)\n",
        "    torch.cuda.manual_seed_all(seed)\n",
        "    torch.backends.cudnn.deterministic = True\n",
        "    np.random.seed(seed)\n",
        "    random.seed(seed)\n",
        "\n",
        "set_seed(cfg.seed) # set seed, default is 1\n",
        "if cfg.save_path is not None:  # make save_path dir\n",
        "    if not os.path.exists(cfg.save_path):\n",
        "        os.makedirs(cfg.save_path, exist_ok=True)\n",
        "print(cfg)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "AvHJg3o0wRPo"
      },
      "source": [
        "## Process dataset"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "e-ZTJOBswSlE"
      },
      "outputs": [],
      "source": [
        "class MMPNERProcessor(object):\n",
        "    def __init__(self, data_path, args) -> None:\n",
        "        self.data_path = data_path\n",
        "        self.tokenizer = BertTokenizer.from_pretrained(args.bert_name, do_lower_case=True)\n",
        "        self.clip_processor = CLIPProcessor.from_pretrained(args.vit_name)\n",
        "        self.aux_processor = CLIPProcessor.from_pretrained(args.vit_name)\n",
        "        self.aux_processor.feature_extractor.size, self.aux_processor.feature_extractor.crop_size = args.aux_size, args.aux_size\n",
        "        self.rcnn_processor = CLIPProcessor.from_pretrained(args.vit_name)\n",
        "        self.rcnn_processor.feature_extractor.size, self.rcnn_processor.feature_extractor.crop_size = args.rcnn_size, args.rcnn_size\n",
        "        self.cwd = args.cwd\n",
        "\n",
        "    def load_from_file(self, mode=\"train\"):\n",
        "        load_file = os.path.join(self.cwd,self.data_path[mode])\n",
        "        logger.info(\"Loading data from {}\".format(load_file))\n",
        "        with open(load_file, \"r\", encoding=\"utf-8\") as f:\n",
        "            lines = f.readlines()\n",
        "            raw_words, raw_targets = [], []\n",
        "            raw_word, raw_target = [], []\n",
        "            imgs = []\n",
        "            for line in lines:\n",
        "                if line.startswith(\"IMGID:\"):\n",
        "                    img_id = line.strip().split('IMGID:')[1] + '.jpg'\n",
        "                    imgs.append(img_id)\n",
        "                    continue\n",
        "                if line != \"\\n\":\n",
        "                    raw_word.append(line.split('\\t')[0])\n",
        "                    label = line.split('\\t')[1][:-1]\n",
        "                    if 'OTHER' in label:\n",
        "                        label = label[:2] + 'MISC'\n",
        "                    raw_target.append(label)\n",
        "                else:\n",
        "                    raw_words.append(raw_word)\n",
        "                    raw_targets.append(raw_target)\n",
        "                    raw_word, raw_target = [], []\n",
        "\n",
        "        assert len(raw_words) == len(raw_targets) == len(imgs), \"{}, {}, {}\".format(len(raw_words), len(raw_targets), len(imgs))\n",
        "        aux_imgs = None\n",
        "        # if not self.use_clip_vit:\n",
        "        aux_path = os.path.join(self.cwd,self.data_path[mode+\"_auximgs\"])\n",
        "        aux_imgs = torch.load(aux_path)\n",
        "\n",
        "        rcnn_imgs = torch.load(os.path.join(self.cwd,self.data_path['train_auximgs']))\n",
        "\n",
        "        return {\"words\": raw_words, \"targets\": raw_targets, \"imgs\": imgs, \"aux_imgs\":aux_imgs, \"rcnn_imgs\":rcnn_imgs}\n",
        "\n",
        "\n",
        "class MMPNERDataset(Dataset):\n",
        "    def __init__(self, processor, label_mapping, img_path=None, aux_path=None, rcnn_img_path=None, max_seq=40, ignore_idx=-100, aux_size=128, rcnn_size=64, mode='train',cwd='') -> None:\n",
        "        self.processor = processor\n",
        "        self.data_dict = processor.load_from_file(mode)\n",
        "        self.tokenizer = processor.tokenizer\n",
        "        self.label_mapping = label_mapping\n",
        "        self.max_seq = max_seq\n",
        "        self.ignore_idx = ignore_idx\n",
        "        self.img_path = img_path\n",
        "        self.aux_img_path = aux_path[mode]  if aux_path is not None else None\n",
        "        self.rcnn_img_path = rcnn_img_path\n",
        "        self.mode = mode\n",
        "        self.clip_processor = self.processor.clip_processor\n",
        "        self.aux_processor = self.processor.aux_processor\n",
        "        self.rcnn_processor = self.processor.rcnn_processor\n",
        "        self.aux_size = aux_size\n",
        "        self.rcnn_size = rcnn_size\n",
        "        self.cwd = cwd\n",
        "    \n",
        "    def __len__(self):\n",
        "        return len(self.data_dict['words'])\n",
        "\n",
        "    def __getitem__(self, idx):\n",
        "        word_list, label_list, img = self.data_dict['words'][idx], self.data_dict['targets'][idx], self.data_dict['imgs'][idx]\n",
        "        tokens, labels = [], []\n",
        "        for i, word in enumerate(word_list):\n",
        "            token = self.tokenizer.tokenize(word)\n",
        "            tokens.extend(token)\n",
        "            label = label_list[i]\n",
        "            for m in range(len(token)):\n",
        "                if m == 0:\n",
        "                    labels.append(self.label_mapping[label])\n",
        "                else:\n",
        "                    labels.append(self.label_mapping[\"X\"])\n",
        "        if len(tokens) >= self.max_seq - 1:\n",
        "            tokens = tokens[0:(self.max_seq - 2)]\n",
        "            labels = labels[0:(self.max_seq - 2)]\n",
        "\n",
        "        encode_dict = self.tokenizer.encode_plus(tokens, max_length=self.max_seq, truncation=True, padding='max_length')\n",
        "        input_ids, token_type_ids, attention_mask = encode_dict['input_ids'], encode_dict['token_type_ids'], encode_dict['attention_mask']\n",
        "        labels = [self.label_mapping[\"[CLS]\"]] + labels + [self.label_mapping[\"[SEP]\"]] + [self.ignore_idx]*(self.max_seq-len(labels)-2)\n",
        "\n",
        "        \n",
        "        if self.img_path is not None:\n",
        "            self.img_path = os.path.join(self.cwd,self.img_path)\n",
        "            # image process\n",
        "            try:\n",
        "                img_path = os.path.join(self.img_path, img)\n",
        "                image = Image.open(img_path).convert('RGB')\n",
        "                image = self.clip_processor(images=image, return_tensors='pt')['pixel_values'].squeeze()\n",
        "            except:\n",
        "                img_path = os.path.join(self.img_path, 'inf.png')\n",
        "                image = Image.open(img_path).convert('RGB')\n",
        "                image = self.clip_processor(images=image, return_tensors='pt')['pixel_values'].squeeze()\n",
        "\n",
        "            if self.aux_img_path is not None:\n",
        "                aux_imgs = []\n",
        "                aux_img_paths = []\n",
        "                if img in self.data_dict['aux_imgs']:\n",
        "                    aux_img_paths  = self.data_dict['aux_imgs'][img]\n",
        "                    aux_img_paths = [os.path.join(self.aux_img_path, path) for path in aux_img_paths]\n",
        "                for i in range(min(3, len(aux_img_paths))):\n",
        "                    aux_img = Image.open(os.path.join(self.cwd,aux_img_paths[i])).convert('RGB')\n",
        "                    aux_img = self.aux_processor(images=aux_img, return_tensors='pt')['pixel_values'].squeeze()\n",
        "                    aux_imgs.append(aux_img)\n",
        "\n",
        "                for i in range(3-len(aux_imgs)):\n",
        "                    aux_imgs.append(torch.zeros((3, self.aux_size, self.aux_size))) \n",
        "\n",
        "                aux_imgs = torch.stack(aux_imgs, dim=0)\n",
        "                assert len(aux_imgs) == 3\n",
        "\n",
        "                if self.rcnn_img_path is not None:\n",
        "                    rcnn_imgs = []\n",
        "                    rcnn_img_paths = []\n",
        "                    img = img.split('.')[0]\n",
        "                    if img in self.data_dict['rcnn_imgs']:\n",
        "                        rcnn_img_paths = self.data_dict['rcnn_imgs'][img]\n",
        "                        rcnn_img_paths = [os.path.join(self.rcnn_img_path, path) for path in rcnn_img_paths]\n",
        "                    for i in range(min(3, len(rcnn_img_paths))):\n",
        "                        rcnn_img = Image.open(rcnn_img_paths[i]).convert('RGB')\n",
        "                        rcnn_img = self.rcnn_processor(images=rcnn_img, return_tensors='pt')['pixel_values'].squeeze()\n",
        "                        rcnn_imgs.append(rcnn_img)\n",
        "\n",
        "                    for i in range(3-len(rcnn_imgs)):\n",
        "                        rcnn_imgs.append(torch.zeros((3, self.rcnn_size, self.rcnn_size))) \n",
        "\n",
        "                    rcnn_imgs = torch.stack(rcnn_imgs, dim=0)\n",
        "                    assert len(rcnn_imgs) == 3\n",
        "                    return torch.tensor(input_ids), torch.tensor(token_type_ids), torch.tensor(attention_mask), torch.tensor(labels), image, aux_imgs, rcnn_imgs\n",
        "\n",
        "                return torch.tensor(input_ids), torch.tensor(token_type_ids), torch.tensor(attention_mask), torch.tensor(labels), image, aux_imgs\n",
        "\n",
        "        assert len(input_ids) == len(token_type_ids) == len(attention_mask) == len(labels)\n",
        "        return torch.tensor(input_ids), torch.tensor(token_type_ids), torch.tensor(attention_mask), torch.tensor(labels)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ZNrxEpTtwoLz"
      },
      "outputs": [],
      "source": [
        "label_mapping = {label:idx for idx, label in enumerate(LABEL_LIST, 1)}\n",
        "label_mapping[\"PAD\"] = 0\n",
        "data_path, img_path, aux_path = DATA_PATH[cfg.dataset_name], IMG_PATH[cfg.dataset_name], AUX_PATH[cfg.dataset_name]\n",
        "rcnn_img_path = DATA_PATH[cfg.dataset_name]['rcnn_img_path']\n",
        "\n",
        "processor = MMPNERProcessor(data_path, cfg)\n",
        "train_dataset = MMPNERDataset(processor, label_mapping, img_path, aux_path, rcnn_img_path, max_seq=cfg.max_seq, ignore_idx=cfg.ignore_idx, aux_size=cfg.aux_size, rcnn_size=cfg.rcnn_size, mode='train',cwd=cwd)\n",
        "train_dataloader = DataLoader(train_dataset, batch_size=cfg.batch_size, shuffle=True, num_workers=4, pin_memory=True)\n",
        "\n",
        "dev_dataset = MMPNERDataset(processor, label_mapping, img_path, aux_path, rcnn_img_path, max_seq=cfg.max_seq, ignore_idx=cfg.ignore_idx, aux_size=cfg.aux_size, rcnn_size=cfg.rcnn_size, mode='dev',cwd=cwd)\n",
        "dev_dataloader = DataLoader(dev_dataset, batch_size=cfg.batch_size, shuffle=False, num_workers=4, pin_memory=True)\n",
        "\n",
        "test_dataset = MMPNERDataset(processor, label_mapping, img_path, aux_path, rcnn_img_path, max_seq=cfg.max_seq, ignore_idx=cfg.ignore_idx, aux_size=cfg.aux_size, rcnn_size=cfg.rcnn_size, mode='test',cwd=cwd)\n",
        "test_dataloader = DataLoader(test_dataset, batch_size=cfg.batch_size, shuffle=False, num_workers=4, pin_memory=True)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "NwrbWDNJw0J4"
      },
      "source": [
        "## Model construction"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Fdw-NO-6w10u"
      },
      "outputs": [],
      "source": [
        "import torch\n",
        "from torch import nn\n",
        "from torchcrf import CRF\n",
        "import torch.nn.functional as F\n",
        "from .modeling_IFA import IFAModel\n",
        "from .clip.modeling_clip import CLIPModel\n",
        "from .clip.configuration_clip import CLIPConfig\n",
        "from transformers import BertConfig, BertModel\n",
        "\n",
        "class IFANERCRFModel(nn.Module):\n",
        "    def __init__(self, label_list, args):\n",
        "        super(IFANERCRFModel, self).__init__()\n",
        "        self.args = args\n",
        "        self.vision_config = CLIPConfig.from_pretrained(self.args.vit_name).vision_config\n",
        "        self.text_config = BertConfig.from_pretrained(self.args.bert_name)\n",
        "\n",
        "        clip_model_dict = CLIPModel.from_pretrained(self.args.vit_name).vision_model.state_dict()\n",
        "        bert_model_dict = BertModel.from_pretrained(self.args.bert_name).state_dict()\n",
        "\n",
        "        print(self.vision_config)\n",
        "        print(self.text_config)\n",
        "\n",
        "        self.vision_config.device = args.device\n",
        "        self.model = IFAModel(self.vision_config, self.text_config)\n",
        "\n",
        "        self.num_labels  = len(label_list) + 1  # pad\n",
        "        self.crf = CRF(self.num_labels, batch_first=True)\n",
        "        self.fc = nn.Linear(self.text_config.hidden_size, self.num_labels)\n",
        "        self.dropout = nn.Dropout(0.1)\n",
        "\n",
        "        # load:\n",
        "        vision_names, text_names = [], []\n",
        "        model_dict = self.model.state_dict()\n",
        "        for name in model_dict:\n",
        "            if 'vision' in name:\n",
        "                clip_name = name.replace('vision_', '').replace('model.', '')\n",
        "                if clip_name in clip_model_dict:\n",
        "                    vision_names.append(clip_name)\n",
        "                    model_dict[name] = clip_model_dict[clip_name]\n",
        "            elif 'text' in name:\n",
        "                text_name = name.replace('text_', '').replace('model.', '')\n",
        "                if text_name in bert_model_dict:\n",
        "                    text_names.append(text_name)\n",
        "                    model_dict[name] = bert_model_dict[text_name]\n",
        "        assert len(vision_names) == len(clip_model_dict) and len(text_names) == len(bert_model_dict), \\\n",
        "                    (len(vision_names), len(text_names), len(clip_model_dict), len(bert_model_dict))\n",
        "        self.model.load_state_dict(model_dict)\n",
        "\n",
        "    def forward(\n",
        "            self, \n",
        "            input_ids=None, \n",
        "            attention_mask=None, \n",
        "            token_type_ids=None, \n",
        "            labels=None, \n",
        "            images=None, \n",
        "            aux_imgs=None,\n",
        "            rcnn_imgs=None,\n",
        "    ):\n",
        "        bsz = input_ids.size(0)\n",
        "\n",
        "        output = self.model(input_ids=input_ids,\n",
        "                            attention_mask=attention_mask,\n",
        "                            token_type_ids=token_type_ids,\n",
        "\n",
        "                            pixel_values=images,\n",
        "                            aux_values=aux_imgs, \n",
        "                            rcnn_values=rcnn_imgs,\n",
        "                            return_dict=True,)\n",
        "\n",
        "        sequence_output = output.last_hidden_state       # bsz, len, hidden\n",
        "        sequence_output = self.dropout(sequence_output)  # bsz, len, hidden\n",
        "        emissions = self.fc(sequence_output)             # bsz, len, labels\n",
        "        \n",
        "        logits = self.crf.decode(emissions, attention_mask.byte())\n",
        "        loss = None\n",
        "        if labels is not None:\n",
        "            loss = -1 * self.crf(emissions, labels, mask=attention_mask.byte(), reduction='mean')  # 去掉CLS\n",
        "            return logits, loss\n",
        "        return logits, None"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "-er4F-4Pw98v"
      },
      "outputs": [],
      "source": [
        "model = IFANERCRFModel(LABEL_LIST, cfg)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "whhxJQVzxBfg"
      },
      "source": [
        "## Train the model"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "JYX7Np0txDWJ"
      },
      "outputs": [],
      "source": [
        "class Trainer(object):\n",
        "    def __init__(self, train_data=None, dev_data=None, test_data=None, model=None, process=None, label_map=None, args=None, logger=None,  writer=None) -> None:\n",
        "        self.train_data = train_data\n",
        "        self.dev_data = dev_data\n",
        "        self.test_data = test_data\n",
        "        self.model = model\n",
        "        self.process = process\n",
        "        self.logger = logger\n",
        "        self.label_map = label_map\n",
        "        self.writer = writer\n",
        "        self.refresh_step = 2\n",
        "        self.best_dev_metric = 0\n",
        "        self.best_test_metric = 0\n",
        "        self.best_dev_epoch = None\n",
        "        self.best_test_epoch = None\n",
        "        self.optimizer = None\n",
        "        self.step = 0\n",
        "        self.args = args\n",
        "        if self.train_data is not None:\n",
        "            self.train_num_steps = len(self.train_data) * args.num_epochs\n",
        "            self.multiModal_before_train()\n",
        "        \n",
        "        \n",
        "    \n",
        "    def train(self):\n",
        "        self.step = 0\n",
        "        self.model.train()\n",
        "        self.logger.info(\"***** Running training *****\")\n",
        "        self.logger.info(\"  Num instance = %d\", len(self.train_data)*self.args.batch_size)\n",
        "        self.logger.info(\"  Num epoch = %d\", self.args.num_epochs)\n",
        "        self.logger.info(\"  Batch size = %d\", self.args.batch_size)\n",
        "        self.logger.info(\"  Learning rate = {}\".format(self.args.lr))\n",
        "        self.logger.info(\"  Evaluate begin = %d\", self.args.eval_begin_epoch)\n",
        "\n",
        "        if self.args.load_path is not None:  # load model from load_path\n",
        "            self.logger.info(\"Loading model from {}\".format(self.args.load_path))\n",
        "            self.model.load_state_dict(torch.load(self.args.load_path))\n",
        "            self.logger.info(\"Load model successful!\")\n",
        "            \n",
        "        with tqdm(total=self.train_num_steps, postfix='loss:{0:<6.5f}', leave=False, dynamic_ncols=True, initial=self.step) as pbar:\n",
        "            self.pbar = pbar\n",
        "            avg_loss = 0\n",
        "            for epoch in range(1, self.args.num_epochs+1):\n",
        "                y_true, y_pred = [], []\n",
        "                y_true_idx, y_pred_idx = [], []\n",
        "                pbar.set_description_str(desc=\"Epoch {}/{}\".format(epoch, self.args.num_epochs))\n",
        "                for batch in self.train_data:\n",
        "                    self.step += 1\n",
        "                    batch = (tup.to(self.args.device)  if isinstance(tup, torch.Tensor) else tup for tup in batch)\n",
        "                    attention_mask, labels, logits, loss = self._step(batch, mode=\"train\")\n",
        "                    avg_loss += loss.detach().cpu().item()\n",
        "\n",
        "                    loss.backward()\n",
        "                    self.optimizer.step()\n",
        "                    self.scheduler.step()\n",
        "\n",
        "                    self.optimizer.zero_grad()\n",
        "\n",
        "                    if isinstance(logits, torch.Tensor): \n",
        "                        logits = logits.argmax(-1).detach().cpu().numpy()  # batch, seq, 1\n",
        "                    label_ids = labels.to('cpu').numpy()\n",
        "                    input_mask = attention_mask.to('cpu').numpy()\n",
        "                    label_map = {idx:label for label, idx in self.label_map.items()}\n",
        "                    for i, mask in enumerate(input_mask):\n",
        "                        temp_1 = []\n",
        "                        temp_2 = []\n",
        "                        temp_1_idx, temp_2_idx = [], []\n",
        "                        for j, m in enumerate(mask):\n",
        "                            if j == 0:\n",
        "                                continue\n",
        "                            if m:\n",
        "                                if label_map[label_ids[i][j]] != \"X\" and label_map[label_ids[i][j]] != \"[SEP]\":\n",
        "                                    temp_1.append(label_map[label_ids[i][j]])\n",
        "                                    temp_2.append(label_map[logits[i][j]])\n",
        "                                    temp_1_idx.append(label_ids[i][j])\n",
        "                                    temp_2_idx.append(logits[i][j])\n",
        "                            else:\n",
        "                                break\n",
        "                        y_true.append(temp_1)\n",
        "                        y_pred.append(temp_2)\n",
        "                        y_true_idx.append(temp_1_idx)\n",
        "                        y_pred_idx.append(temp_2_idx)\n",
        "\n",
        "                    if self.step % self.refresh_step == 0:\n",
        "                        avg_loss = float(avg_loss) / self.refresh_step\n",
        "                        print_output = \"loss:{:<6.5f}\".format(avg_loss)\n",
        "                        pbar.update(self.refresh_step)\n",
        "                        pbar.set_postfix_str(print_output)\n",
        "                        if self.writer:\n",
        "                            self.writer.log({'avg_loss': avg_loss})\n",
        "                        avg_loss = 0\n",
        "               \n",
        "                if epoch >= self.args.eval_begin_epoch:\n",
        "                    if self.dev_data:\n",
        "                        self.evaluate(epoch)   # generator to dev.\n",
        "                    if self.test_data:\n",
        "                        self.test(epoch)\n",
        "            \n",
        "            torch.cuda.empty_cache()\n",
        "            \n",
        "            pbar.close()\n",
        "            self.pbar = None\n",
        "            self.logger.info(\"Get best dev performance at epoch {}, best dev f1 score is {}\".format(self.best_dev_epoch, self.best_dev_metric))\n",
        "            self.logger.info(\"Get best test performance at epoch {}, best test f1 score is {}\".format(self.best_test_epoch, self.best_test_metric))\n",
        "\n",
        "    def evaluate(self, epoch):\n",
        "        self.model.eval()\n",
        "        self.logger.info(\"***** Running evaluate *****\")\n",
        "        self.logger.info(\"  Num instance = %d\", len(self.dev_data)*self.args.batch_size)\n",
        "        self.logger.info(\"  Batch size = %d\", self.args.batch_size)\n",
        "\n",
        "        y_true, y_pred = [], []\n",
        "        y_true_idx, y_pred_idx = [], []\n",
        "        step = 0\n",
        "        with torch.no_grad():\n",
        "            with tqdm(total=len(self.dev_data), leave=False, dynamic_ncols=True) as pbar:\n",
        "                pbar.set_description_str(desc=\"Dev\")\n",
        "                total_loss = 0\n",
        "                for batch in self.dev_data:\n",
        "                    step += 1\n",
        "                    batch = (tup.to(self.args.device)  if isinstance(tup, torch.Tensor) else tup for tup in batch)  # to cpu/cuda device\n",
        "                    attention_mask, labels, logits, loss = self._step(batch, mode=\"dev\")    # logits: batch, seq, num_labels\n",
        "                    total_loss += loss.detach().cpu().item()\n",
        "\n",
        "                    if isinstance(logits, torch.Tensor):    \n",
        "                        logits = logits.argmax(-1).detach().cpu().numpy()  # batch, seq, 1\n",
        "                    label_ids = labels.detach().cpu().numpy()\n",
        "                    input_mask = attention_mask.detach().cpu().numpy()\n",
        "                    label_map = {idx:label for label, idx in self.label_map.items()}\n",
        "                    for i, mask in enumerate(input_mask):\n",
        "                        temp_1 = []\n",
        "                        temp_2 = []\n",
        "                        temp_1_idx, temp_2_idx = [], []\n",
        "                        for j, m in enumerate(mask):\n",
        "                            if j == 0:\n",
        "                                continue\n",
        "                            if m:\n",
        "                                if label_map[label_ids[i][j]] != \"X\" and label_map[label_ids[i][j]] != \"[SEP]\":\n",
        "                                    temp_1.append(label_map[label_ids[i][j]])\n",
        "                                    temp_2.append(label_map[logits[i][j]])\n",
        "                                    temp_1_idx.append(label_ids[i][j])\n",
        "                                    temp_2_idx.append(logits[i][j])\n",
        "                            else:\n",
        "                                break\n",
        "                        y_true.append(temp_1)\n",
        "                        y_pred.append(temp_2)\n",
        "                        y_true_idx.append(temp_1_idx)\n",
        "                        y_pred_idx.append(temp_2_idx)\n",
        "                    \n",
        "                    pbar.update()\n",
        "                # evaluate done\n",
        "                pbar.close()\n",
        "\n",
        "                results = classification_report(y_true, y_pred, digits=4)  \n",
        "                self.logger.info(\"***** Dev Eval results *****\")\n",
        "                self.logger.info(\"\\n%s\", results)\n",
        "                f1_score = float(results.split('\\n')[-4].split('      ')[-2].split('    ')[-1])\n",
        "                if self.writer: \n",
        "                    self.writer.log({'eva_f1': f1_score})\n",
        "\n",
        "                self.logger.info(\"Epoch {}/{}, best dev f1: {}, best epoch: {}, current dev f1 score: {}.\"\\\n",
        "                            .format(epoch, self.args.num_epochs, self.best_dev_metric, self.best_dev_epoch, f1_score))\n",
        "                if f1_score >= self.best_dev_metric:  # this epoch get best performance\n",
        "                    self.logger.info(\"Get better performance at epoch {}\".format(epoch))\n",
        "                    self.best_dev_epoch = epoch\n",
        "                    self.best_dev_metric = f1_score # update best metric(f1 score)\n",
        "                    if self.args.save_path is not None: # save model\n",
        "                        torch.save(self.model.state_dict(), self.args.save_path+\"/best_model.pth\")\n",
        "                        self.logger.info(\"Save best model at {}\".format(self.args.save_path))\n",
        "               \n",
        "\n",
        "        self.model.train()\n",
        "\n",
        "    def test(self, epoch):\n",
        "        self.model.eval()\n",
        "        self.logger.info(\"\\n***** Running testing *****\")\n",
        "        self.logger.info(\"  Num instance = %d\", len(self.test_data)*self.args.batch_size)\n",
        "        self.logger.info(\"  Batch size = %d\", self.args.batch_size)\n",
        "\n",
        "        if self.args.load_path is not None:  # load model from load_path\n",
        "            self.logger.info(\"Loading model from {}\".format(self.args.load_path))\n",
        "            self.model.load_state_dict(torch.load(self.args.load_path))\n",
        "            self.logger.info(\"Load model successful!\")\n",
        "        y_true, y_pred = [], []\n",
        "        y_true_idx, y_pred_idx = [], []\n",
        "        with torch.no_grad():\n",
        "            with tqdm(total=len(self.test_data), leave=False, dynamic_ncols=True) as pbar:\n",
        "                pbar.set_description_str(desc=\"Testing\")\n",
        "                for batch in self.test_data:\n",
        "                    batch = (tup.to(self.args.device)  if isinstance(tup, torch.Tensor) else tup for tup in batch)  # to cpu/cuda device\n",
        "                    attention_mask, labels, logits, loss = self._step(batch, mode=\"dev\")    # logits: batch, seq, num_labels\n",
        "            \n",
        "                    if isinstance(logits, torch.Tensor):    #\n",
        "                        logits = logits.argmax(-1).detach().cpu().tolist()  # batch, seq, 1\n",
        "                    label_ids = labels.detach().cpu().numpy()\n",
        "                    input_mask = attention_mask.detach().cpu().numpy()\n",
        "                    label_map = {idx:label for label, idx in self.label_map.items()}\n",
        "                    for i, mask in enumerate(input_mask):\n",
        "                        temp_1 = []\n",
        "                        temp_2 = []\n",
        "                        temp_1_idx, temp_2_idx = [], []\n",
        "                        for j, m in enumerate(mask):\n",
        "                            if j == 0:\n",
        "                                continue\n",
        "                            if m:\n",
        "                                if label_map[label_ids[i][j]] != \"X\" and label_map[label_ids[i][j]] != \"[SEP]\":\n",
        "                                    temp_1.append(label_map[label_ids[i][j]])\n",
        "                                    temp_2.append(label_map[logits[i][j]])\n",
        "                                    temp_1_idx.append(label_ids[i][j])\n",
        "                                    temp_2_idx.append(logits[i][j])\n",
        "                            else:\n",
        "                                break\n",
        "                        y_true.append(temp_1)\n",
        "                        y_pred.append(temp_2)\n",
        "                        y_true_idx.append(temp_1_idx)\n",
        "                        y_pred_idx.append(temp_2_idx)\n",
        "                    \n",
        "                    pbar.update()\n",
        "                # evaluate done\n",
        "                pbar.close()\n",
        "\n",
        "                results = classification_report(y_true, y_pred, digits=4) \n",
        "                self.logger.info(\"***** Test Eval results *****\")\n",
        "                self.logger.info(\"\\n%s\", results)\n",
        "                f1_score = float(results.split('\\n')[-4].split('      ')[-2].split('    ')[-1])\n",
        "                if self.writer:\n",
        "                    self.writer.log({'test_f1': f1_score})\n",
        "                total_loss = 0\n",
        "                \n",
        "                self.logger.info(\"Epoch {}/{}, best test f1: {}, best epoch: {}, current test f1 score: {}.\"\\\n",
        "                            .format(epoch, self.args.num_epochs, self.best_test_metric, self.best_test_epoch, f1_score))\n",
        "                if f1_score >= self.best_test_metric:  # this epoch get best performance\n",
        "                    self.best_test_metric = f1_score\n",
        "                    self.best_test_epoch = epoch\n",
        "                   \n",
        "        self.model.train()\n",
        "\n",
        "\n",
        "    def predict(self):\n",
        "        self.model.eval()\n",
        "        self.logger.info(\"\\n***** Running predicting *****\")\n",
        "        self.logger.info(\"  Num instance = %d\", len(self.test_data)*self.args.batch_size)\n",
        "        self.logger.info(\"  Batch size = %d\", self.args.batch_size)\n",
        "        if self.args.load_path is not None:  # load model from load_path\n",
        "            self.logger.info(\"Loading model from {}\".format(self.args.load_path))\n",
        "            self.model.load_state_dict(torch.load(self.args.load_path))\n",
        "            self.logger.info(\"Load model successful!\")\n",
        "            self.model.to(self.args.device)\n",
        "        y_pred = []\n",
        "\n",
        "        with torch.no_grad():\n",
        "            with tqdm(total=len(self.test_data), leave=False, dynamic_ncols=True) as pbar:\n",
        "                pbar.set_description_str(desc=\"Predicting\")\n",
        "                for batch in self.test_data:\n",
        "                    batch = (tup.to(self.args.device)  if isinstance(tup, torch.Tensor) else tup for tup in batch)  # to cpu/cuda device\n",
        "                    attention_mask, labels, logits, loss = self._step(batch, mode=\"dev\")    # logits: batch, seq, num_labels\n",
        "            \n",
        "                    if isinstance(logits, torch.Tensor):    # \n",
        "                        logits = logits.argmax(-1).detach().cpu().tolist()  # batch, seq, 1\n",
        "                    label_ids = labels.detach().cpu().numpy()\n",
        "                    input_mask = attention_mask.detach().cpu().numpy()\n",
        "                    label_map = {idx:label for label, idx in self.label_map.items()}\n",
        "                    for i, mask in enumerate(input_mask):\n",
        "                        temp_1 = []\n",
        "                        for j, m in enumerate(mask):\n",
        "                            if j == 0:\n",
        "                                continue\n",
        "                            if m:\n",
        "                                if label_map[label_ids[i][j]] != \"X\" and label_map[label_ids[i][j]] != \"[SEP]\":\n",
        "                                    temp_1.append(label_map[logits[i][j]])\n",
        "                            else:\n",
        "                                break\n",
        "                        y_pred.append(temp_1)\n",
        "                    \n",
        "                    pbar.update()\n",
        "                # evaluate done\n",
        "                pbar.close()\n",
        "        \n",
        "    def _step(self, batch, mode=\"train\"):\n",
        "        input_ids, token_type_ids, attention_mask, labels, images, aux_imgs, rcnn_imgs = batch\n",
        "        logits, loss = self.model(input_ids=input_ids, attention_mask=attention_mask, token_type_ids=token_type_ids, labels=labels, images=images, aux_imgs=aux_imgs, rcnn_imgs=rcnn_imgs)\n",
        "        return attention_mask, labels, logits, loss\n",
        "\n",
        "\n",
        "    def multiModal_before_train(self):\n",
        "        # bert lr\n",
        "        parameters = []\n",
        "        params = {'lr':self.args.lr, 'weight_decay':1e-2}\n",
        "        params['params'] = []\n",
        "        for name, param in self.model.named_parameters():\n",
        "            if 'text' in name:\n",
        "                params['params'].append(param)\n",
        "        parameters.append(params)\n",
        "\n",
        "         # vit lr\n",
        "        params = {'lr':3e-5, 'weight_decay':1e-2}\n",
        "        params['params'] = []\n",
        "        for name, param in self.model.named_parameters():\n",
        "            if 'vision' in name:\n",
        "                params['params'].append(param)\n",
        "        parameters.append(params)\n",
        "\n",
        "        # crf lr\n",
        "        params = {'lr':5e-2, 'weight_decay':1e-2}\n",
        "        params['params'] = []\n",
        "        for name, param in self.model.named_parameters():\n",
        "            if 'crf' in name or name.startswith('fc'):\n",
        "                params['params'].append(param)\n",
        "        parameters.append(params)\n",
        "\n",
        "        self.optimizer = optim.AdamW(parameters)\n",
        "\n",
        "        self.model.to(self.args.device)\n",
        "            \n",
        "        self.scheduler = get_linear_schedule_with_warmup(optimizer=self.optimizer, \n",
        "                                 num_warmup_steps=self.args.warmup_ratio*self.train_num_steps, \n",
        "                                 num_training_steps=self.train_num_steps)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ZteN3icdxNRY"
      },
      "outputs": [],
      "source": [
        "trainer = Trainer(train_data=train_dataloader, dev_data=dev_dataloader, test_data=test_dataloader, model=model, label_map=label_mapping, args=cfg, logger=logger, writer=writer)\n",
        "trainer.train()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "V9t5HxUpxQmQ"
      },
      "source": [
        "## Predict"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "iO3xWuD4xRiC"
      },
      "outputs": [],
      "source": [
        "trainer = Trainer(train_data=None, dev_data=None, test_data=test_dataloader, model=model, label_map=label_mapping, args=cfg, logger=logger, writer=writer)\n",
        "trainer.predict()"
      ]
    }
  ],
  "metadata": {
    "accelerator": "GPU",
    "colab": {
      "name": "multimodal_ner_tutorial.ipynb",
      "provenance": []
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}
